/*
 * Copyright 2010 Bo Zhou<Bo.Schwarzstein@gmail.com>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 
		http://www.apache.org/licenses/LICENSE-2.0 
 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "MeshRelaxationDeformer.h"

#include "MeshSmoothing.h"
#include "BilateralMeshSmoothing.h"
#include "LaplacianMeshSmoothing.h"
#include "TaubinMeshSmoothing.h"

#include <assert.h>
#include <memory>
#include <set>
#include <valarray>

#include <maya/MDataHandle.h>
#include <maya/MFloatPointArray.h>
#include <maya/MFnEnumAttribute.h>
#include <maya/MFnMesh.h>
#include <maya/MFnMeshData.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnPlugin.h>
#include <maya/MFnTypedAttribute.h>
#include <maya/MGlobal.h>
#include <maya/MItGeometry.h>
#include <maya/MPointArray.h>
#include <maya/MProgressWindow.h>

void* MeshRelaxationDeformer::Creator()
{
	return new MeshRelaxationDeformer();
}

MStatus MeshRelaxationDeformer::Initialize()
{
	MFnEnumAttribute EnumAttrBuilder;
	MFnNumericAttribute NumAttrBuilder;
	MFnTypedAttribute TypedAttrBuilder;

	MStatus R(MStatus::kFailure);

	// AttrMethod
	//
	AttrMethod = EnumAttrBuilder.create("relaxationMethod","rxm",0);
	EnumAttrBuilder.addField("Average",0);
	EnumAttrBuilder.addField("Laplacian",1);
	EnumAttrBuilder.addField("Taubin",2);
	EnumAttrBuilder.addField("Bilateral",3);
	EnumAttrBuilder.setKeyable(false);
	CHECK_MSTATUS_AND_RETURN( addAttribute(AttrMethod), MStatus::kFailure );

	// AttrTime
	//
	AttrTime = NumAttrBuilder.create("relaxationTime","rxt",MFnNumericData::kInt,0);
	NumAttrBuilder.setKeyable(false);
	CHECK_MSTATUS_AND_RETURN(addAttribute(AttrTime), MStatus::kFailure);

	// AttrShowProgress
	//
	AttrShowProgress = NumAttrBuilder.create("relaxationShowProgress","rxsp",MFnNumericData::kBoolean,0,&R);
	NumAttrBuilder.setKeyable(false);
	R = addAttribute(AttrShowProgress);
	CHECK_MSTATUS_AND_RETURN(R, MStatus::kFailure);

	// AttrUsePaintedVertices
	//
	AttrUsePaintedVertices = NumAttrBuilder.create("relaxationUsePaintedVertices","rxupv",MFnNumericData::kBoolean,0,&R);
	NumAttrBuilder.setKeyable(false);
	R = addAttribute(AttrUsePaintedVertices);
	CHECK_MSTATUS_AND_RETURN(R, MStatus::kFailure);

	// AttrPaintedColor
	//
	AttrPaintedColor = NumAttrBuilder.createColor("relaxationPaintedColor","rxpc",&R);
	NumAttrBuilder.setKeyable(false);
	NumAttrBuilder.setConnectable(false);
	R = addAttribute(AttrPaintedColor);
	CHECK_MSTATUS_AND_RETURN(R, MStatus::kFailure);

	// AttrLambda
	//
	AttrLambda = NumAttrBuilder.create("relaxationLambda","rxl",MFnNumericData::kFloat);
	NumAttrBuilder.setDefault(0.1f);
	NumAttrBuilder.setMin(0.001f);
	NumAttrBuilder.setMax(0.999f);
	NumAttrBuilder.setKeyable(false);
	CHECK_MSTATUS_AND_RETURN(addAttribute(AttrLambda), MStatus::kFailure);

	// AttrMu
	//
	AttrMu = NumAttrBuilder.create("relaxationMu","rxu",MFnNumericData::kFloat);
	NumAttrBuilder.setDefault(0.1f);
	NumAttrBuilder.setMin(0.001f);
	NumAttrBuilder.setMax(0.999f);
	NumAttrBuilder.setKeyable(false);
	CHECK_MSTATUS_AND_RETURN(addAttribute(AttrMu), MStatus::kFailure);

	// AttrWeightMethod
	//
	AttrWeightMethod = EnumAttrBuilder.create("relaxationWeightMethod","rxwm");
	EnumAttrBuilder.addField("Equal",0);
	EnumAttrBuilder.addField("Fujiwara",1);
	NumAttrBuilder.setDefault(0);
	NumAttrBuilder.setKeyable(false);
	CHECK_MSTATUS_AND_RETURN(addAttribute(AttrWeightMethod), MStatus::kFailure);

	// AttrSigmaC
	//
	AttrSigmaC = NumAttrBuilder.create("relaxationSigmaC","rxsigc",MFnNumericData::kFloat);
	NumAttrBuilder.setDefault(0.1f);
	NumAttrBuilder.setKeyable(false);
	CHECK_MSTATUS_AND_RETURN(addAttribute(AttrSigmaC), MStatus::kFailure);

	// AttrSigmaS
	//
	AttrSigmaS = NumAttrBuilder.create("relaxationSigmaS","rxsigs",MFnNumericData::kFloat);
	NumAttrBuilder.setDefault(0.1f);
	NumAttrBuilder.setKeyable(false);
	CHECK_MSTATUS_AND_RETURN(addAttribute(AttrSigmaS), MStatus::kFailure);

	// Setup relationship between attributes and outputGeom
	//
	attributeAffects(AttrMethod, outputGeom);
	attributeAffects(AttrTime, outputGeom);
	attributeAffects(AttrShowProgress, outputGeom);
	attributeAffects(AttrUsePaintedVertices, outputGeom);
	attributeAffects(AttrPaintedColor, outputGeom);
	
	attributeAffects(AttrLambda, outputGeom);
	attributeAffects(AttrMu, outputGeom);
	attributeAffects(AttrWeightMethod, outputGeom);

	attributeAffects(AttrSigmaC, outputGeom);
	attributeAffects(AttrSigmaS, outputGeom);

	return MStatus::kSuccess;
}

MTypeId MeshRelaxationDeformer::Identifier(0x98A3BD);

MObject MeshRelaxationDeformer::AttrMethod;
MObject MeshRelaxationDeformer::AttrTime;
MObject MeshRelaxationDeformer::AttrShowProgress;
MObject MeshRelaxationDeformer::AttrUsePaintedVertices;
MObject MeshRelaxationDeformer::AttrPaintedColor;

MObject MeshRelaxationDeformer::AttrLambda;
MObject MeshRelaxationDeformer::AttrMu;
MObject MeshRelaxationDeformer::AttrWeightMethod;

MObject MeshRelaxationDeformer::AttrSigmaC;
MObject MeshRelaxationDeformer::AttrSigmaS;

MeshRelaxationDeformer::MeshRelaxationDeformer()
{
}

MeshRelaxationDeformer::~MeshRelaxationDeformer()
{
}

MStatus MeshRelaxationDeformer::compute(const MPlug& plug, MDataBlock& dataBlock)
{
	MStatus R = MStatus::kUnknownParameter;

	if( plug.attribute() == outputGeom )
	{
		// Prepare Input Attributes.
		//
		unsigned int LogicalIndex = plug.logicalIndex();
		MPlug InPlug(thisMObject(),input);
		InPlug.selectAncestorLogicalIndex(LogicalIndex,input);

		MDataHandle DHInput = dataBlock.inputValue(InPlug,&R);
		MDataHandle DHInputGeom = DHInput.child(inputGeom);
		MDataHandle DHGroup = DHInput.child(groupId);

		MDataHandle DHMethod = dataBlock.inputValue(AttrMethod);
		MDataHandle DHTime = dataBlock.inputValue(AttrTime);
		MDataHandle DHShowProgress = dataBlock.inputValue(AttrShowProgress);
		MDataHandle DHUsePaintedVertices = dataBlock.inputValue(AttrUsePaintedVertices);
		MDataHandle DHPaintedColor = dataBlock.inputValue(AttrPaintedColor);

		MDataHandle DHLambda = dataBlock.inputValue(AttrLambda);
		MDataHandle DHMu = dataBlock.inputValue(AttrMu);
		MDataHandle DHWeightMethod = dataBlock.inputValue(AttrWeightMethod);

		MDataHandle DHSigmaC =  dataBlock.inputValue(AttrSigmaC);
		MDataHandle DHSigmaS =  dataBlock.inputValue(AttrSigmaS);

		MFnData::Type GeomType = DHInputGeom.type();
		if( GeomType != MFnData::kMesh )
		{
			MGlobal::displayError("Can only deform kMesh!");
			return MStatus::kFailure;
		}
		// Output attributes.
		//
		MDataHandle DHOutputGeom = dataBlock.outputValue(plug,&R);
		CHECK_MSTATUS_AND_RETURN(R,MStatus::kFailure);
		DHOutputGeom.copy(DHInputGeom);

		// Copy base mesh as output geomery.
		//
		MFnMeshData InputMeshData(DHInputGeom.data(),&R);
		CHECK_MSTATUS_AND_RETURN(R,MStatus::kFailure);

		MFnMeshData OutputMeshData(DHOutputGeom.data(),&R);
		CHECK_MSTATUS_AND_RETURN(R,MStatus::kFailure);

		// Calculate !
		//
		std::auto_ptr<MeshSmoothing> Smoother;

		const int Method = DHMethod.asInt();
		switch(Method)
		{
		case 1:
			{
				LaplacianMeshSmoothing* p = new LaplacianMeshSmoothing( OutputMeshData.object() );
				p->SetWeightType( (const LaplacianMeshSmoothing::WeightType)DHWeightMethod.asInt() );
				p->SetLambda( DHLambda.asFloat() );
				Smoother.reset(p);
			}
			break;
		case 2:
			{
				TaubinMeshSmoothing* p = new TaubinMeshSmoothing( OutputMeshData.object() );
				p->SetWeightType( (const LaplacianMeshSmoothing::WeightType)DHWeightMethod.asInt() );
				p->SetLambda( DHLambda.asFloat() );
				p->SetMu( DHMu.asFloat() );
				Smoother.reset(p);
			}
			break;
		case 3:
			{
				BilateralMeshSmoothing* p = new BilateralMeshSmoothing( OutputMeshData.object() );
				p->SetSigma( DHSigmaC.asFloat(), DHSigmaS.asFloat() );
				Smoother.reset(p);
			}
			break;
		default:
			{
				MeshSmoothing* p = new MeshSmoothing( OutputMeshData.object());
				Smoother.reset(p);
			}
		}

		// If smoothing on specificated facets.
		//
		if( DHUsePaintedVertices.asBool() )
		{
			Smoother->UsePaintedVertices();

			const float3& PC = DHPaintedColor.asFloat3();
			Smoother->SetPaintedColor( MColor(PC) );
		}

		bool PreparedOk = Smoother->Prepare();
		if(!PreparedOk)
		{
			DHOutputGeom.setClean();
			return MStatus::kFailure;
		}

		const bool ShowProgress = DHShowProgress.asBool();
		const int Time = DHTime.asInt();
		if(ShowProgress)
		{
			MProgressWindow::reserve();
			MProgressWindow::setTitle("Relaxing Mesh...");
			MProgressWindow::setInterruptable(true);
			MProgressWindow::setProgressRange(0,Time);
		}

		if(ShowProgress)
		{
			MProgressWindow::startProgress();
		}

		for(int i=0; i<Time; ++i)
		{
			if(ShowProgress)
			{
				if( MProgressWindow::isCancelled() )
					break;
				MProgressWindow::setProgress(i);
			}
			Smoother->Smooth();
		}

		if(ShowProgress)
		{
			MProgressWindow::endProgress();
		}

		DHOutputGeom.setClean();
	}

	return R;
}

void MeshRelaxationDeformer::setModifiedCallback(MSelectionList & list, bool listAdded)
{
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#ifdef _WIN32
	#define API __declspec(dllexport)
#else
	#define API
#endif

MStatus API initializePlugin( MObject O )
{
	MStatus R;
	MFnPlugin Plugin(O,"Bo Zhou <bo.schwarzstein@gmail.com>","beta 2","9.0");
	R = Plugin.registerNode( "meshRelaxation", MeshRelaxationDeformer::Identifier, MeshRelaxationDeformer::Creator, MeshRelaxationDeformer::Initialize, MPxNode::kDeformerNode );
	return R;
}

MStatus API uninitializePlugin( MObject O)
{
	MStatus R;
	MFnPlugin Plugin(O);
	R = Plugin.deregisterNode( MeshRelaxationDeformer::Identifier );
	return R;
}

